/*
    Copyright 2007-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Definitions and explicit instantiations of the ir_grid class
    members. */

#include <CCfits/CCfits>
#include "ir_grid.h"
#include "grid.h"
#include "arepo_grid.h"
#include "arepo_emission.h"
#include "blitz-fits.h"
#include "fits-utilities.h"
#include "mpi_util.h"
#include "boost/lexical_cast.hpp"

#ifdef WITH_AREPO
extern "C" {
#include "allvars.h"
}
#endif

using namespace blitz;

/** Constructor loads cell data from a FITS file.  The grid structure
    is determined by the supplied T_grid_impl object, but to get the
    length unit we still need the structure_hdu. */
template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
ir_grid (boost::shared_ptr<T_grid_impl> g,
	 CCfits::ExtHDU& structure_hdu,
	 CCfits::ExtHDU& data_hdu,
	 const density_generator& dg,
	 const T_content& lambda) :
  emission_collection<polychromatic_policy, sampling_policy, rng_policy>(),
  g_(g), lambda_(independent_copy(lambda)), L_lambda_(contiguousArray), 
  L_lambda_old_(contiguousArray), m_dust_(contiguousArray),
  tot_lum_(0)
{
  structure_hdu.readKey("lengthunit", units_ ["length"]);

  // get the amount of dust in cells.  For this we need the density generator
  const int nc = this->n_cells();
  m_dust_.resize(nc, dg.make_density_vector(0,0).size());

  // read amount of metals in cells
  array_1 m_gas, m_metal;
  read (data_hdu.column("mass_gas"), m_gas);
  read (data_hdu.column("mass_metals"), m_metal);

  // loop over cells and create dust mass array and emitter objects
  int i = 0;
  for (iterator c = this->begin (), e = this->end(); 
       c != e; ++c, ++i) {
    // density generator takes and returns DENSITIES, but we want MASS
    m_dust_(i, Range::all()) = 
      dg.make_density_vector(m_gas(i)/c.volume(), m_metal(i)/c.volume()) *
      c.volume();
  }

  make_emitters();
}

// helper function to call the arepo_emission method if that's the type
template<typename T> void update_bounding_boxes(const T&) {assert(0);};
template<typename T> void 
update_bounding_boxes(const mcrx::arepo_grid<T>& g) {
  T::T_emitter::update_bounding_boxes(g); };


/** Constructor getting the data from the Arepo data structures. */
template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
ir_grid (boost::shared_ptr<T_grid_impl> g,
	 const density_generator& dg,
	 const T_content& lambda) :
  emission_collection<polychromatic_policy, sampling_policy, rng_policy>(),
  g_(g), lambda_(independent_copy(lambda)), L_lambda_(contiguousArray), 
  L_lambda_old_(contiguousArray), m_dust_(contiguousArray),
  tot_lum_(0)
{
#ifdef WITH_AREPO

  // First copy units from the arepo units
  units_["length"]=arepo::arepo_units.get("length");
  units_["time"]=arepo::arepo_units.get("time");
  units_["mass"]=arepo::arepo_units.get("mass");

  // to convert arepo density to our units. (convert call kind of
  // unnecessary since we know the units are the same from above, but
  // done for clarity)
  const T_float convert_density = 
    arepo::mcon/(arepo::lcon*arepo::lcon*arepo::lcon)*
    units::convert(arepo::arepo_units.get("mass")+"/"+
		   arepo::arepo_units.get("length")+"^3",
		   units_.get("mass")+"/"+units_.get("length")+"^3");

  // get the amount of dust in cells.  For this we need the density generator
  const int nc = this->n_cells();
  m_dust_.resize(nc, dg.make_density_vector(0,0).size());

  int i = 0;
  for (iterator c = this->begin (), e = this->end(); 
       c != e; ++c, ++i) {

    assert(c.volume()>0);
    const T_float rho_gas = SphP[i].Density*convert_density;
    const T_float rho_met = rho_gas*SphP[i].Metallicity;

    // density generator takes and returns DENSITIES, but we want MASS
    m_dust_(i, Range::all()) = 
      dg.make_density_vector(rho_gas, rho_met) * c.volume();
  }

  make_emitters();
  // since this is the arepo grid, we also now need to set up the
  // bounding boxes of the emitters
  update_bounding_boxes(*g_);

#else
  cerr << "Not compiled with Arepo support" << endl;
  exit(1);
#endif
}

/** The ir_grid is not like a conventional grid in that the quantities
    are stored in the grid cells. Because the main function of this
    grid is to calculate dust emission, those quantities are stored in
    the grid object as an array, not in the cells. This means that we
    are unable to build this grid the traditional way using a grid
    factory. Instead this constructor uses a structure vector and a
    mass vector. */
template <typename grid_type,
	  template<typename> class sampling_policy, 
	  typename rng_policy>
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
ir_grid (boost::shared_ptr<T_grid_impl> g,
	 const array_2& mass,
	 const T_content& lambda,
	 const T_unit_map& units) :
  emission_collection<polychromatic_policy, sampling_policy, rng_policy>(),
  g_(g), m_dust_(mass), lambda_(independent_copy(lambda)),
  L_lambda_(contiguousArray), 
  L_lambda_old_(contiguousArray), tot_lum_(0)
{
  units_["length"] = units.get("length");
  // Because the dust mass is already supplied, we just need to set up
  // the emitters.
  make_emitters();
}


template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
void
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
make_emitters()
{
  const int nc = this->n_cells();
  const int n_lambda = lambda_.size();

  // Allocate SED array
  cout << "Allocating a (" << nc << ", " << n_lambda
       << ") memory block for cell emission, "
       << 1.0*nc*n_lambda*sizeof(T_float)/(1024*1024*1024) << " GB.\n";
  L_lambda_.resize(nc, n_lambda);
  // Set to invalid value
  L_lambda_ = blitz::quiet_NaN(T_float());

  int i = 0;
  array_1 temp_lum;
  for (iterator c = this->begin (), e = this->end(); 
       c != e; ++c, ++i) {

    // if the cell has no data member, default-construct one
    if(!c->data()) {
      c->set_data(new T_data() );
    }

    // create an emitter object referencing the main L_lambda array
    temp_lum.weakReference(L_lambda_(i, Range::all ()));
    assert (temp_lum.isStorageContiguous());
    
    T_emitter dd (temp_lum, &lambda_, c);
    c->data()->get_emitter() = dd;
  }
}



template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
CCfits::ExtHDU*
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
create_hdu(CCfits::FITS& file, const TinyVector<int,2>& size, 
	   const string& name,
	   const string& cmt,
	   const string& unit,
	   const string& unitcmt) const
{
  std::vector<long> naxes;
  naxes.push_back(size[0]);
  naxes.push_back(size[1]);
  
  // check if intensity HDU already exists
  CCfits::ExtHDU* hdu;    
  try {
    hdu = &open_HDU (file, name);
  }
  catch (CCfits::FITS::NoSuchHDU&) {
    cout << "Creating " << name << " HDU with size " 
	 << size << endl;
    
    hdu = file.addImage (name, DOUBLE_IMG, naxes);
    file.setCompressionType(0);
    
    hdu->writeComment(cmt);

    hdu->addKey("imunit", unit, unitcmt);
  }
  return hdu;
}


template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
void
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
write_seds(CCfits::FITS& file, const string& hdu_name,
	   bool normalized_for_sampling)
{
  const string cmt("This HDU contains the dust emission SEDs in the grid cells. The first dimension is grid cell, the second wavelength.");
  const string unit("W/m");
  const string unitcmt("Cell emission in units of "+unit);

  if(false/*this->shared_domain()*/) {
    // shared domain. coadd arrays.
    mpi_sum_arrays(L_lambda_);
    if(is_mpi_master()) {
      std::cout << "Writing cell SEDs" << std::endl;
      CCfits::ExtHDU* hdu = create_hdu(file, L_lambda_.shape(), hdu_name,
				       cmt, unit, unitcmt);
      write(*hdu, L_lambda_);
      hdu->writeChecksum();
    }
  }
  else {    // split domain. Need to write consecutively
    {
      if(normalized_for_sampling) {
	// undo normalization so we have the actual luminosities
	L_lambda_ = L_lambda_*L_emitted_(tensor::i)*(1./prev_total_weight_);
      }
      array_2 out(mpi_stack_arrays(L_lambda_, firstDim));
      if(is_mpi_master()) {
	std::cout << "Writing intensities and energy deposition" << std::endl;
	CCfits::ExtHDU* hdu = create_hdu(file, out.shape(), hdu_name,
					 cmt, unit, unitcmt);
	write(*hdu, out);
	hdu->writeChecksum();
      }
    }
  }
}

template <typename grid_type, 
	  template<typename> class sampling_policy, 
	  typename rng_policy>
void
mcrx::ir_grid<grid_type, sampling_policy, rng_policy>::
write_temps(CCfits::FITS& file, const string& hdu_name)
{
  const string cmt("This HDU contains the dust temperatures in the grid cells.  The first dimension is grid cell, the second grain species.");
  const string unit("K");
  const string unitcmt("Dust temperature in units of "+unit);

  if(is_mpi_master())
    std::cout << "Writing dust temperatures" << std::endl;

  for(int i=0; i<temps_.size(); ++i) {
    array_2 out(mpi_stack_arrays(temps_[i], firstDim));
    if(is_mpi_master()) {
      const string name = hdu_name + boost::lexical_cast<string>(i);
      CCfits::ExtHDU* hdu = create_hdu(file, out.shape(), name,
				       cmt, unit, unitcmt);
      write(*hdu, out);
      hdu->writeChecksum();
    }
  }
}

// Explicit instantiations for the two cases we use. If someone starts
// using rejection sampling or invents a new grid, these need to be
// added here.
namespace mcrx {
  template class ir_grid<
    adaptive_grid<
      cell_data<
	typename emitter_types<adaptive_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > >,
    cumulative_sampling, local_random >;
#ifdef WITH_AREPO
  template class ir_grid<
    arepo_grid<
      cell_data<
	typename emitter_types<arepo_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > >,
    cumulative_sampling, local_random>;
#endif
}
